"""
摄像头实时检测（类别统计）GUI
--------------------------------
1. 通过本地摄像头采集画面
2. 将 JPEG 帧 POST 到远端推理服务器 /detect
3. 服务器返回：
   - cls_list : 检测到的类别列表
   - img      : 画好框的 BGR 图像 hex 字符串
4. 主线程负责把结果显示在 PyQt5 界面
   - 左侧：原图
   - 右侧：推理结果图 + 类别统计
5. 支持 开始 / 暂停 / 继续 / 返回主界面
"""

import sys
import cv2
import requests
import numpy as np
from collections import Counter
from PyQt5 import QtCore, QtGui, QtWidgets
from PyQt5.QtCore import QThread
from PyQt5.QtWidgets import QSizePolicy, QMessageBox

# 读取远端服务器地址
from main import load_settings

cfg = load_settings()
SERVER_URL = f"http://{cfg['ip']}:{cfg['port']}/detect"


# =============================================================================
# 后台推理线程
# =============================================================================
class InferenceThread(QThread):
    """
    继承 QThread，在后台完成：
    1. 摄像头采集
    2. 网络推理
    3. 把结果通过信号抛回主线程
    """
    # 自定义信号：rgb, ann_rgb, cls_text
    frame_ready = QtCore.pyqtSignal(np.ndarray, np.ndarray, str)

    def __init__(self, url, parent=None):
        super().__init__(parent)
        self.url = url
        self._running = True  # 总开关
        self._pause = False  # 暂停开关
        self.session = requests.Session()  # 复用 TCP 连接
        self.cap = None  # cv2.VideoCapture 实例

    # -------------------------------------------------------------------------
    # 线程主循环
    # -------------------------------------------------------------------------
    def run(self):
        # 打开默认摄像头（0）
        self.cap = cv2.VideoCapture(0)
        if not self.cap.isOpened():
            return

        # 设置分辨率
        self.cap.set(cv2.CAP_PROP_FRAME_WIDTH, 640)
        self.cap.set(cv2.CAP_PROP_FRAME_HEIGHT, 480)

        while self._running:
            # 如果处于暂停状态，简单 sleep 后继续
            if self._pause:
                self.msleep(50)
                continue

            # 1. 读取一帧
            ret, frame = self.cap.read()
            if not ret:
                break

            # 2. 编码为 JPEG
            rgb = cv2.cvtColor(frame, cv2.COLOR_BGR2RGB)  # 用于本地显示
            ok, enc = cv2.imencode('.jpg', frame)
            if not ok:
                continue

            # 3. 发送到服务器
            try:
                resp = self.session.post(
                        self.url,
                        files={'file': ('frame.jpg', enc.tobytes(), 'image/jpeg')},
                        timeout=5
                        )
                if resp.status_code == 200:
                    data = resp.json()
                    cls_list = data.get("cls_list", [])
                    img_hex = data.get("img", "")

                    # 解码服务器返回的带框图像
                    img_arr = np.frombuffer(bytes.fromhex(img_hex), np.uint8)
                    ann = cv2.imdecode(img_arr, cv2.IMREAD_COLOR)
                    ann_rgb = cv2.cvtColor(ann, cv2.COLOR_BGR2RGB)

                    # 类别统计
                    counter = Counter(cls_list)
                    stat = "\n".join([f"{k}：{v}" for k, v in counter.items()])
                    cls_text = stat if stat else "未检测到物体"
                else:
                    cls_text = "网络错误"
                    ann_rgb = rgb
            except Exception as e:
                cls_text = "推理错误"
                ann_rgb = rgb

            # 4. 通过信号把结果抛回主线程
            self.frame_ready.emit(rgb, ann_rgb, cls_text)

            # 5. 控制帧率，避免 GUI 卡死
            self.msleep(30)

        # 释放摄像头
        self.cap.release()

    # -------------------------------------------------------------------------
    # 对外接口
    # -------------------------------------------------------------------------
    def pause(self):
        """暂停"""
        self._pause = True

    def resume(self):
        """继续"""
        self._pause = False

    def stop(self):
        """完全停止线程"""
        self._running = False
        self.terminate()  # 强制终止，不阻塞
        self.wait(500)  # 最多等 500ms


# =============================================================================
# 主窗口
# =============================================================================
class MainWindow(QtWidgets.QMainWindow):

    def __init__(self):
        super().__init__()
        self.setWindowTitle("摄像头实时检测（类别统计）")
        self.resize(1700, 1000)
        self.setMaximumSize(1700, 1000)  # 防止窗口被拉大
        self.setStyleSheet("background-color: #fffef8;")

        # -------------------- 中央布局 --------------------
        central = QtWidgets.QWidget(self)
        self.setCentralWidget(central)
        main_layout = QtWidgets.QHBoxLayout(central)
        main_layout.setContentsMargins(10, 10, 10, 10)

        # -------------------- 左侧 --------------------
        left_w = QtWidgets.QWidget()
        left_w.setMaximumWidth(650)
        left_w.setSizePolicy(QSizePolicy.Fixed, QSizePolicy.Expanding)
        left_layout = QtWidgets.QVBoxLayout(left_w)

        # 原图标签
        self.label_src = QtWidgets.QLabel("原图")
        self.label_src.setMinimumHeight(200)
        self.label_src.setAlignment(QtCore.Qt.AlignCenter)
        self.label_src.setStyleSheet(
                """
                            background-color: #ffffff;
                            color: #5d8a6a;
                            font-family: "PingFang SC", sans-serif;
                            font-size: 20px;
                            border: 1px solid #5d8a6a;
                            border-radius: 8px;
                            padding: 10px;
                        """
                )
        # 防止图像被拉伸变形
        self.label_src.setScaledContents(True)
        self.label_src.setSizePolicy(
                QtWidgets.QSizePolicy.Ignored,
                QtWidgets.QSizePolicy.Ignored
                )
        left_layout.addWidget(self.label_src, 2)

        # 开始 / 暂停 / 返回 按钮
        self.btn_start = QtWidgets.QPushButton("开始检测")
        self.btn_start.clicked.connect(self.start_camera)
        left_layout.addWidget(self.btn_start)

        self.btn_pause = QtWidgets.QPushButton("暂停")
        self.btn_pause.clicked.connect(self.toggle_pause)
        left_layout.addWidget(self.btn_pause)

        self.btn_back = QtWidgets.QPushButton("返回")
        self.btn_back.clicked.connect(self.go_to_main_1)
        left_layout.addWidget(self.btn_back)

        # -------------------- 右侧 --------------------
        right_w = QtWidgets.QWidget()
        right_layout = QtWidgets.QVBoxLayout(right_w)

        # 检测结果图
        self.label_dst = QtWidgets.QLabel("检测结果")
        self.label_dst.setAlignment(QtCore.Qt.AlignCenter)
        self.label_dst.setStyleSheet(
                """
                            background-color: #f2f5f9;
                            color: #5d8a6a;
                            font-family: "PingFang SC", sans-serif;
                            font-size: 36px;
                            border: 2px dashed #5d8a6a;
                            border-radius: 10px;
                        """
                )
        self.label_dst.setScaledContents(True)
        self.label_dst.setSizePolicy(
                QtWidgets.QSizePolicy.Ignored,
                QtWidgets.QSizePolicy.Ignored
                )
        right_layout.addWidget(self.label_dst, 2)

        # 类别统计标签
        self.label_class = QtWidgets.QLabel("类别统计：-")
        self.label_class.setAlignment(QtCore.Qt.AlignLeft | QtCore.Qt.AlignTop)
        self.label_class.setWordWrap(True)
        self.label_class.setStyleSheet(
                """
                            background-color:#fff;
                            border:1px solid #5d8a6a;
                            border-radius:8px;
                            padding:10px;
                        """
                )
        right_layout.addWidget(self.label_class, 1)

        # 把左右两部分塞进主布局
        main_layout.addWidget(left_w, 1)
        main_layout.addWidget(right_w, 2)

        # -------------------- 线程实例 --------------------
        self.thread = None

    # -------------------------------------------------------------------------
    # 槽函数
    # -------------------------------------------------------------------------
    def start_camera(self):
        """点击【开始检测】时触发"""
        if self.thread is None or not self.thread.isRunning():
            self.thread = InferenceThread(SERVER_URL)
            self.thread.frame_ready.connect(self.on_frame)
            self.thread.start()

    def toggle_pause(self):
        """暂停 / 继续 切换"""
        if self.thread:
            if self.thread._pause:
                self.thread.resume()
                self.btn_pause.setText("暂停")
            else:
                self.thread.pause()
                self.btn_pause.setText("继续")

    def on_frame(self, rgb, ann_rgb, cls_text):
        """线程每推来一帧就调用"""
        self.display(rgb, self.label_src)  # 更新原图
        self.display(ann_rgb, self.label_dst)  # 更新结果图
        self.label_class.setText(cls_text)  # 更新类别统计

    def display(self, rgb, label):
        """把 numpy RGB 图像显示到 QLabel"""
        h, w, ch = rgb.shape
        # 构造 QImage（注意传入 data 生命周期）
        qimg = QtGui.QImage(rgb.data, w, h, w * ch, QtGui.QImage.Format_RGB888)
        label.setPixmap(
                QtGui.QPixmap.fromImage(qimg).scaled(
                        label.size(),
                        QtCore.Qt.KeepAspectRatio,
                        QtCore.Qt.SmoothTransformation
                        )
                )

    # -------------------------------------------------------------------------
    # 关闭 & 返回
    # -------------------------------------------------------------------------
    def closeEvent(self, event):
        """窗口关闭时确保线程退出"""
        if self.thread and self.thread.isRunning():
            self.thread.stop()
        event.accept()

    def go_to_main_1(self):
        """返回主界面（main.py 里的 UI_main_1）"""
        try:
            from main import UI_main_1

            self.win = UI_main_1()
            self.win.show()
            self.close()
        except ImportError:
            QMessageBox.warning(self, "提示", "找不到 main.py，无法返回主界面！")


# =============================================================================
# 程序入口
# =============================================================================
if __name__ == "__main__":
    app = QtWidgets.QApplication(sys.argv)
    win = MainWindow()
    win.show()
    sys.exit(app.exec_())
